import os
import json
import numpy as np
from glob import glob
import cv2
from tqdm import tqdm
import base64


def polygon2json(bboxes, categorys, image_path, json_dir):
    '''

        :param bbox: [[[x1,y1,x2,y2,x3,y3]]]
        :param category: [['dog','cat']]
        :param score: [[0.9,0.8]]
        :param image_path: '/home/i.jpg'
        :return:
        '''
    assert len(bboxes) == len(categorys)
    assert len(image_path) == len(bboxes)
    
    for i in range(len(image_path)):
        res = {"version": "4.2.10",
               "flags": {}, }
        json_path = os.path.join(json_dir, image_path[i].split('/')[-1].split('.')[0] + '.json')
        with open(json_path, encoding='utf8', mode='w') as f:
            res['imagePath'] = os.path.basename(image_path[i])
            bbox = bboxes[i]
            category = categorys[i]
            img = cv2.imread(image_path[i])
            h, w, c = img.shape
            imag = cv2.imencode('.jpg', img)[1]
            base64_data = str(base64.b64encode(imag))[2:-1]
            # b64=base64.b64encode(img).decode('utf-8')
            
            cur_res = []
            # print(category)
            for j in range(len(bbox)):
                cur_poly = bbox[j]
                cur_res.append(
                    {'label': category[j], 'points': [[poly[0], poly[1]] for poly in cur_poly], "group_id": 'null',
                     "shape_type": "rectangle", "flags": {}})
            res['shapes'] = cur_res
            # res['time_labeled']=13141516
            # res['labeled']='true'
            # im = cv2.imread(image_path[i])
            # h, w, c = im.shape
            res['imageHeight'] = h
            res['imageWidth'] = w
            res['imageData'] = base64_data
            res = json.dumps(res)
            f.write(res)
def IoU(bbox, gt):
    """
    :param bbox: (n, 4)
    :param gt: (m, 4)
    :return: (n, m)
    numpy 广播机制 从后向前对齐。 维度为1 的可以重复等价为任意维度
    eg: (4,3,2)   (3,2)  (3,2)会扩充为(4,3,2)
        (4,1,2)   (3,2) (4,1,2) 扩充为(4, 3, 2)  (3, 2)扩充为(4, 3,2) 扩充的方法为重复
    广播会在numpy的函数 如sum, maximun等函数中进行
    pytorch同理。
    扩充维度的方法：
    eg: a  a.shape: (3,2)  a[:, None, :] a.shape: (3, 1, 2) None 对应的维度相当于newaxis
    """
    lt = np.maximum(bbox[:, None, :2], gt[:, :2])  # left_top (x, y)
    rb = np.minimum(bbox[:, None, 2:], gt[:, 2:])  # right_bottom (x, y)
    wh = np.maximum(rb - lt + 1, 0)                # inter_area (w, h)
    inter_areas = wh[:, :, 0] * wh[:, :, 1]        # shape: (n, m)
    box_areas = (bbox[:, 2] - bbox[:, 0] + 1) * (bbox[:, 3] - bbox[:, 1] + 1)
    gt_areas = (gt[:, 2] - gt[:, 0] + 1) * (gt[:, 3] - gt[:, 1] + 1)
    IoU = inter_areas / (box_areas[:, None] + gt_areas - inter_areas)
    return IoU

def read_labelme(json_file):
    data = open(json_file, "r")
    data = json.load(data)
    shapes = data["shapes"]
    bboxes = []
    for shape in shapes:
        pt = np.array(shape['points'], dtype=np.int32)
        bboxes.append(pt)
    bboxes = np.array(bboxes)
    bboxes = bboxes.reshape(-1, 2, 2)
    return bboxes,data["imagePath"]

def get_obj(img_dir,json_dir):
    files=glob(json_dir+'/*.json')
    obj_list=[]
    for file in tqdm(files):
        bboxes,im_name=read_labelme(file)
        im=cv2.imread(os.path.join(img_dir,im_name))
        for bb in bboxes:
            cur_obj=im[bb[0][1]:bb[1][1],bb[0][0]:bb[1][0],:]
            obj_list.append(cur_obj)
    return obj_list

def blend_img(bg,fg):
    assert bg.shape[0]==fg.shape[0] and bg.shape[1]==fg.shape[1]
    # alpha=0.1
    # bg=bg*alpha+fg*(1-alpha)
    # return bg.astype(np.uint8)
    if np.random.randint(0,100)>50:
        fg=np.fliplr(fg)
    src_mask = 255 * np.ones(fg.shape[0:2], fg.dtype)
    # 位置 此处图片不能偏移出去，大小控制好
    center = (fg.shape[1]//2,fg.shape[0]//2)
    # Clone seamlessly.  提供了两种方式 cv2.MIXED_CLONE 和 cv2.NORMAL_CLONE 结果不同的
    #mode=cv2.NORMAL_CLONE if np.random.randint(0,100)>50 else cv2.NORMAL_CLONE
    output = cv2.seamlessClone(fg, bg, src_mask, center, cv2.NORMAL_CLONE)
    return output

def render_bg(bg,fg_list,thresh=0.1,scales=[0.5,1.2]):
    H,W,C=bg.shape
    bbox=[]
    for fg in fg_list:
        scale=np.random.uniform(scales[0],scales[1])
        h,w,c=fg.shape
        fg=cv2.resize(fg,(int(w*scale),int(h*scale)))
        h, w, c = fg.shape
        x,y=np.random.randint(0,W-w-1),np.random.randint(0,H-h-1)
        if len(bbox)==0:
            bbox.append([x,y,x+w,y+h])
            bg[y:y+h,x:x+w,:]=blend_img(bg[y:y+h,x:x+w,:],fg)
        else:
            iter=0
            while True:
                if iter>100:
                    break
                iou=IoU(np.array([[x,y,x+w,y+h]]),np.array(bbox))
                iter+=1
                if iou.max()<thresh:
                    bbox.append([x, y, x + w, y + h])
                    bg[y:y + h, x:x + w, :] = blend_img(bg[y:y + h, x:x + w, :], fg)
                    break
                else:
                    x, y = np.random.randint(0, W - w - 1), np.random.randint(0, H - h - 1)
    return bg,bbox
if __name__=='__main__':
    fg_paths=['/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/1y',
              '/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/2y','/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/211013',
              '/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/211013_1','/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/y'
              ]
    bg_paths='/media/wsl/SB@data/公开数据集/无人机/bg'
    bg_files=glob(bg_paths+'/*.png')
    obj_list=[]
    for fg in fg_paths:
        oo=get_obj(fg,fg)
        obj_list.extend(oo)
    obj_list=get_obj('/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/1y','/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/1y')
    
    dst_dir='/media/wsl/SB@data/公开数据集/无人机/已标注/20210915/211015_syn'
    num_obj=(3,6)
    for i in tqdm(range(1000)):
        bg_file=np.random.choice(bg_files)
        bg=cv2.imread(bg_file)
        fg_list=np.random.choice(obj_list,np.random.randint(num_obj[0],num_obj[1]))
        bg,bbox=render_bg(bg,fg_list)
        im_path=os.path.join(dst_dir,'%d.png'%i)
        cv2.imwrite(im_path,bg)
        polygon2json([np.array(bbox).reshape(-1,2,2).tolist()],[['0']*len(bbox)],[im_path],dst_dir)
        # cv2.imshow('a',bg)
        # cv2.waitKey(0)